package org.leiyang.common.LoadBalancer;

import org.springframework.aop.support.AopUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
import org.springframework.stereotype.Component;

@Component
public class LoadBalancerClientPostProcesser implements BeanPostProcessor, BeanFactoryAware {

    private LoadBalancerClientFactory loadBalancerClientFactory;

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        Class<?> beanType = AopUtils.getTargetClass(bean);
        if(BlockingLoadBalancerClient.class.isAssignableFrom(beanType)) {
            BlockingLoadBalancerClient loadBalancerClient = (BlockingLoadBalancerClient) bean;
            return proxyLoadBalancerClient(loadBalancerClientFactory, loadBalancerClient);
        }

        return BeanPostProcessor.super.postProcessAfterInitialization(bean, beanName);
    }

    private LoadBalancerClient proxyLoadBalancerClient(LoadBalancerClientFactory loadBalancerClientFactory,
                                                       BlockingLoadBalancerClient blockingLoadBalancerClient) {
        return new CustomBlockingLoadBalancerClient(loadBalancerClientFactory, blockingLoadBalancerClient);
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.loadBalancerClientFactory = beanFactory.getBean(LoadBalancerClientFactory.class);
    }
}
